"""Simplified advection experiments matching the paper exactly."""
import sys
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import time

sys.path.insert(0, str(Path(__file__).parent.parent))

from src.reproduct_paper.problems.advection import advection_1d_reference
from src.reproduct_paper.advection_solver import AdvectionSolver


def run_figure3a_experiment():
    """Figure 3(a): Effect of number of sub-domains."""
    print("=" * 80)
    print("Figure 3(a): Effect of number of sub-domains")
    print("Fixed: Q=20x20, M=300, Rm=1.0, t_f=10, Nb=10")
    print("=" * 80)

    problem = advection_1d_reference()

    # Test configurations
    configs = [
        {"Ne": 2, "Nx": 2, "Nt": 1, "label": "Ne=2 (2x1)"},
        {"Ne": 4, "Nx": 2, "Nt": 2, "label": "Ne=4 (2x2)"},
        {"Ne": 8, "Nx": 4, "Nt": 2, "label": "Ne=8 (4x2)"},
    ]

    results = []

    for config in configs:
        print(f"\nTesting {config['label']}...")
        start_time = time.time()

        solver = AdvectionSolver(
            x_domain=problem["domain_x"],
            t_final=10.0,
            num_time_blocks=10,
            num_subdomains_x=config["Nx"],
            num_subdomains_t=config["Nt"],
            collocation_x=20,
            collocation_t=20,
            hidden_width=300,
            rand_magnitude=1.0,
            wave_speed=problem["wave_speed"],
            initial_condition=problem["initial_condition"],
            exact_solution=problem["exact_solution"],
            rng=np.random.default_rng(1),
        )

        res = solver.solve()
        errors = solver.get_overall_errors()
        elapsed = time.time() - start_time

        print(f"  Time: {elapsed:.2f}s")
        print(f"  Max error: {errors['max']:.6e}")
        print(f"  RMS error: {errors['rms']:.6e}")

        results.append({
            "Ne": config["Ne"],
            "label": config["label"],
            "max_error": errors["max"],
            "rms_error": errors["rms"],
            "time": elapsed,
        })

    # Save results
    import json
    output_file = Path(__file__).parent.parent / "experiments" / "fig3a_results.json"
    with open(output_file, 'w') as f:
        json.dump(results, f, indent=2)

    print(f"\nResults saved to: {output_file}")
    return results


def run_single_baseline():
    """Run single baseline test from Figure 1."""
    print("=" * 80)
    print("Baseline test (Figure 1 configuration)")
    print("Ne=8 (4x2), Q=20x20, M=300, Rm=1.0, t_f=10, Nb=10")
    print("=" * 80)

    problem = advection_1d_reference()

    solver = AdvectionSolver(
        x_domain=problem["domain_x"],
        t_final=10.0,
        num_time_blocks=10,
        num_subdomains_x=4,
        num_subdomains_t=2,
        collocation_x=20,
        collocation_t=20,
        hidden_width=300,
        rand_magnitude=1.0,
        wave_speed=problem["wave_speed"],
        initial_condition=problem["initial_condition"],
        exact_solution=problem["exact_solution"],
        rng=np.random.default_rng(1),
    )

    print("\nSolving...")
    start = time.time()
    results = solver.solve()
    elapsed = time.time() - start
    errors = solver.get_overall_errors()

    print(f"\nTotal time: {elapsed:.2f}s")
    print(f"Average block time: {np.mean(results['block_times']):.2f}s")
    print(f"Max error: {errors['max']:.6e}")
    print(f"RMS error: {errors['rms']:.6e}")

    # Generate solution plot
    print("\nGenerating plots...")
    x_eval = np.linspace(0, 5, 100)
    t_eval = np.linspace(0, 10, 100)
    u_numerical = solver.evaluate_solution(x_eval, t_eval)
    u_exact = problem["exact_solution"](x_eval, t_eval)

    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    X, T = np.meshgrid(x_eval, t_eval, indexing='ij')

    # Numerical solution
    im1 = axes[0].contourf(X, T, u_numerical, levels=20, cmap='viridis')
    axes[0].set_xlabel('x')
    axes[0].set_ylabel('t')
    axes[0].set_title('locELM Solution')
    plt.colorbar(im1, ax=axes[0])

    # Error
    error = np.abs(u_numerical - u_exact)
    im2 = axes[1].contourf(X, T, error, levels=20, cmap='hot')
    axes[1].set_xlabel('x')
    axes[1].set_ylabel('t')
    axes[1].set_title(f'Absolute Error (max={errors["max"]:.2e})')
    plt.colorbar(im2, ax=axes[1])

    plt.tight_layout()
    output_path = Path(__file__).parent.parent / "experiments" / "advection_baseline.png"
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    print(f"Plot saved to: {output_path}")
    plt.close()

    # Save block errors
    block_errors_data = {
        "block_idx": list(range(len(results["block_errors"]))),
        "max_errors": [e["max"] for e in results["block_errors"]],
        "rms_errors": [e["rms"] for e in results["block_errors"]],
    }

    import json
    error_file = Path(__file__).parent.parent / "experiments" / "baseline_block_errors.json"
    with open(error_file, 'w') as f:
        json.dump(block_errors_data, f, indent=2)

    print(f"Block errors saved to: {error_file}")

    return results, errors


if __name__ == "__main__":
    # Run baseline first
    print("Running baseline experiment...\n")
    baseline_results, baseline_errors = run_single_baseline()

    print("\n\nRunning Figure 3(a) experiment...\n")
    fig3a_results = run_figure3a_experiment()

    print("\n" + "=" * 80)
    print("Experiments completed!")
    print("=" * 80)
